#include <arpa/inet.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <unistd.h>

#include <freertos/FreeRTOS.h>
#include <freertos/event_groups.h>
#include <freertos/queue.h>
#include <freertos/task.h>

#include <cJSON.h>
#include <driver/uart.h>
#include <esp_check.h>
#include <esp_event.h>
#include <esp_log.h>
#include <esp_netif.h>
#include <esp_task_wdt.h>
#include <nvs_flash.h>

#include "wifi_station.h"
#include "uart.h"
#include "utilities.h"

static const in_port_t g_tcp_conn_port = htons(8080);
static const in_port_t g_udp_forward_port = htons(8081);
static QueueHandle_t g_udp_addr_queue;

/**
 * @brief Parse a JSON string and extra message to addr
 *      
 * 
 * @param[in]   str Raw JSON string
 * @param[out]  addr Address to be written
 * @return int 
 *              - 2 Successfully parse and write ip and port to addr
 *              - 1 Successfully parse and write port to addr(no key named "ip" in JSON string)
 *              - 0 Parse failed
 * 
 * @note        This function don't check validity of arguments
 *              The key "ip" is not necessary, so addr->sin_addr needs to be set when return 1
 */
static int address_parse(const char *str, struct sockaddr_in *addr) {
    int ret = 0;
    cJSON *root = cJSON_Parse(str);
    cJSON *ip;
    cJSON *port;
    
    addr->sin_family = AF_INET;
    if (NULL == root) {
        ESP_LOGW(TAG, "Not a JSON string");
    } else if (NULL == (port = cJSON_GetObjectItem(root, "port"))) {
        ESP_LOGW(TAG, "No key named \"port\" in JSON");
    } else if (!cJSON_IsNumber(port)) {
        ESP_LOGW(TAG, "The value of \"port\" is not number");
    } else if (NULL == (ip = cJSON_GetObjectItem(root, "ip"))) {
        ret = 1;
        addr->sin_port = htons(port->valueint);
        ESP_LOGI(TAG, "Write port to addr");
    } else if (!cJSON_IsString(ip)) {
        ESP_LOGW(TAG, "The value of \"ip\" is not string");
    } else if (inet_pton(AF_INET, ip->valuestring, &addr->sin_addr) == 0) {
        ESP_LOGW(TAG, "The value of \"ip\" is not a valid ip string");
    } else {
        ret = 2;
        addr->sin_port = htons(port->valueint);
        ESP_LOGI(TAG, "Write ip and port to addr");
    }

    cJSON_Delete(root);
    return ret;
}

static void udp_forward_task(void *pvParameters) {
    enum { xTicksToDelay = pdMS_TO_TICKS(1) > 1 ? pdMS_TO_TICKS(1) : 1 };
    TickType_t xTicksToWait = portMAX_DELAY;
    char ip_str[INET_ADDRSTRLEN];
    uint8_t buffer[256];
    struct sockaddr_in addr = {
        .sin_family = AF_INET,
        .sin_addr.s_addr = IPADDR_ANY,
        .sin_port = g_udp_forward_port,
    };
    int sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_IP);
    (void)bind(sock, (struct sockaddr *)&addr, sizeof(addr));
    const struct timeval timeout = {.tv_sec = 10, .tv_usec = 0};
    setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));

    while (true) {
        if (xQueueReceive(g_udp_addr_queue, &addr, xTicksToWait)) {
            inet_ntop(AF_INET, &addr.sin_addr, ip_str, sizeof(addr));
            ESP_LOGI(TAG, "Got address %s:%u", ip_str, ntohs(addr.sin_port));
            xTicksToWait = 0;
        }
        int len = uart_read_bytes(UART_NUM_1, buffer, sizeof(buffer), 0);
        if (len > 0) {
            int err = sendto(sock, buffer, len, 0, (struct sockaddr *)&addr, sizeof(addr));
            if (err < 0) {
                ESP_LOGE(TAG, "Error occurred during sending: errno %d", errno);
                // xTicksToWait = portMAX_DELAY;
            }
        }
        vTaskDelay(xTicksToDelay);
    }

    vTaskDelete(NULL);
}

static void tcp_client_handler(void *pvParameters) {
    int sock = (uintptr_t)pvParameters;
    char buffer[128];
    struct sockaddr_in to;
    struct sockaddr from;
    socklen_t fromlen;
    ssize_t len = recvfrom(sock, buffer, sizeof(buffer), 0, &from, &fromlen);
    if (len > 0) {
        buffer[len] = 0;
        int err = address_parse(buffer, &to);
        if (1 == err && sizeof(struct sockaddr_in) == fromlen) {
            to.sin_addr.s_addr = ((struct sockaddr_in *)&from)->sin_addr.s_addr;
            xQueueSend(g_udp_addr_queue, &to, portMAX_DELAY);
        } else if (2 == err) {
            xQueueSend(g_udp_addr_queue, &to, portMAX_DELAY);
        } else {
            ESP_LOGW(TAG, "Invalid JSON string");
            goto err;
        }
        size_t len = snprintf(buffer, sizeof(buffer), "{\"port\": %"PRIu16"}", g_udp_forward_port);
        send(sock, buffer, len, 0);
    }
err:
    close(sock);
    vTaskDelete(NULL);
}

static void tcp_server_task(void *pvParameters) {
    int sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    struct sockaddr_in server_addr = {
        .sin_family = AF_INET,
        .sin_port = g_tcp_conn_port,
        .sin_addr.s_addr = INADDR_ANY,
    };
    int tmp = bind(sock, (struct sockaddr *)&server_addr, sizeof(server_addr));
    (void)tmp;
    listen(sock, 5);

    while (true) {
        struct sockaddr_in client_addr;
        socklen_t addr_len = sizeof(client_addr);
        int client_sock = accept(sock, (struct sockaddr *)&client_addr, &addr_len);
        xTaskCreate(tcp_client_handler,
            "tcp_handler",
            4096,
            (void *)(uintptr_t)client_sock,
            2,
            NULL);
    }
}

void app_main(void) {
    esp_err_t ret = nvs_flash_init();
    if (ret == ESP_ERR_NVS_NO_FREE_PAGES || ret == ESP_ERR_NVS_NEW_VERSION_FOUND) {
        ESP_ERROR_CHECK(nvs_flash_erase());
        ret = nvs_flash_init();
    }
    ESP_ERROR_CHECK(ret);

    uart_initialize(UART_NUM_1, 115200, 19, 18);
    wifi_connect_sta("Xiaomi 14 Pro", "1048794538", WIFI_AUTH_WPA2_PSK);
    start_mdns_service("lidar", "lidar");

    g_udp_addr_queue = xQueueCreate(5, sizeof(struct sockaddr_in));

    xTaskCreate(tcp_server_task,
        "tcp_server",
        4096,
        NULL,
        3,
        NULL);

    xTaskCreate(udp_forward_task,
        "udp_forward",
        4096,
        NULL,
        4,
        NULL);
}
